package test.openai.model;

import com.liuqi.openai.core.client.httpclient.OpenAiHttpClient;
import com.liuqi.openai.data.embedding.Embedding;
import com.liuqi.openai.data.segment.TextSegment;
import com.liuqi.openai.model.OpenAiEmbeddingModel;
import com.liuqi.openai.model.OpenAiModelException;
import com.liuqi.openai.model.embedding.EmbeddingModel;
import com.liuqi.openai.model.output.Response;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import lombok.extern.slf4j.Slf4j;

/**
 * @author liuqi
 * @date 2025/7/20
 **/
@Slf4j
public class TextOpenAiEmbeddingModel {
    public static void main(String[] args) throws OpenAiModelException {
        OpenAiHttpClient openAiClient = OpenAiHttpClient.builder()
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                // -DapiKey="apikey"
                .openAiApiKey(System.getProperty("apiKey"))
                .build();

        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .openAiClient(openAiClient)
                .modelName("text-embedding-v1")
                .build();

        List<TextSegment> textSegments = new ArrayList<>();
        textSegments.add(TextSegment.from("1"));
        textSegments.add(TextSegment.from("2"));
        textSegments.add(TextSegment.from("3"));

        // 单个文本 embed
        Response<Embedding> testEmbed = embeddingModel.embed("123");
        System.out.println("testEmbed: " + testEmbed);

        // 多文本 embed
        Response<List<Embedding>> testsEmbed = embeddingModel.embed(Arrays.asList("1", "2", "3"));
        System.out.println("testsEmbed: " + testsEmbed);

        // TextSegment embed
        Response<Embedding> textSegmentEmbed = embeddingModel.embedTextSegment(TextSegment.from("1"));
        System.out.println("textSegmentEmbed: " + textSegmentEmbed);

        // TextSegment集合 embed
        Response<List<Embedding>> textSegmentsEmbed = embeddingModel.embedTextSegment(textSegments);
        System.out.println("textSegmentEmbed: " + textSegmentsEmbed);

        System.out.println();
    }
}
